# **************************************
# --*-- coding: utf-8 --*--
# @Author  : white
# @FileName: kNN分类.py
# @Time    : 2025-08-14
# **************************************
# https://www.joinquant.com/view/community/detail/bb850ee76d1cae16cc587f29c4439ebd
from sklearn import neighbors
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from pathlib import Path

img = Path("images")
img.mkdir(exist_ok=True)

x1 = np.random.normal(50, 6, 200)
y1 = np.random.normal(5, 0.5, 200)

x2 = np.random.normal(30, 6, 200)
y2 = np.random.normal(4, 0.5, 200)

x3 = np.random.normal(45, 6, 200)
y3 = np.random.normal(2.5, 0.5, 200)

plt.scatter(x1, y1, c='b', marker='s', s=50, alpha=0.8)
plt.scatter(x2, y2, c='r', marker='^', s=50, alpha=0.8)
plt.scatter(x3, y3, c='g', s=50, alpha=0.8)
plt.savefig(img.joinpath("KNN分布图.png"))

# 所有的 x 坐标和 y 坐标放在一起
x_val = np.concat((x1, x2, x3))
y_val = np.concat((y1, y2, y3))
# 求出 x 值的最大差还有 y 值的最大差
x_diff = np.max(x_val) - np.min(x_val)
y_diff = np.max(y_val) - np.min(y_val)
# 将坐标除以这个差以归一化，再将 x 和 y 值两两配对
x_normalized = x_val / x_diff
y_normalized = y_val / y_diff
xy_normalized = np.column_stack((x_normalized, y_normalized))
# 生成相应的分类标签
labels = [1] * 200 + [2] * 200 + [3] * 200
# 生成 sklearn 的最近 k 邻分类功能了。参数中，n_neighbors 设为 30，其他的都使用默认值即可
clf = neighbors.KNeighborsClassifier(30)
# 进行数据拟合 TODO 这里如何进行训练的？不应该将这个空间中的所有坐标点一个一个传进去吗？
clf.fit(xy_normalized, labels)
# 首先，我们想知道 (50,5) 和 (30,3) 两个点附近最近的 5 个样本分别都是什么。
# 坐标别忘了除以 x_diff 和 y_diff 来归一化。
nearests = clf.kneighbors([(50 / x_diff, 5 / y_diff), (30 / x_diff, 3 / y_diff)], 10, False)
# 预测
prediction = clf.predict([(50 / x_diff, 5 / y_diff), (30 / x_diff, 3 / y_diff)])
# 预测概率
prediction_proba = clf.predict_proba([(50 / x_diff, 5 / y_diff), (30 / x_diff, 3 / y_diff)])
# (50, 5) 有 100% 的可能性是 1 类，而 (30,3) 有 80% 是 2 类，20% 是3类。
print("预测概率：", prediction_proba)
# 准确率打分
x1_test = np.random.normal(50, 6, 100)
y1_test = np.random.normal(5, 0.5, 100)

x2_test = np.random.normal(30, 6, 100)
y2_test = np.random.normal(4, 0.5, 100)

x3_test = np.random.normal(45, 6, 100)
y3_test = np.random.normal(2.5, 0.5, 100)

xy_test_normalized = np.column_stack((np.concatenate((x1_test, x2_test, x3_test)) / x_diff,
                                      np.concatenate((y1_test, y2_test, y3_test)) / y_diff))
labels_test = [1] * 100 + [2] * 100 + [3] * 100
score = clf.score(xy_test_normalized, labels_test)
print("模型评估：", score)

# 生成穷举空间坐标点
xx, yy = np.meshgrid(np.arange(1, 70.1, 0.1), np.arange(1, 7.01, 0.01))
# 坐标点归一化
xx_normalized = xx / x_diff
yy_normalized = yy / y_diff
# 特征堆叠
coords = np.column_stack((xx_normalized.ravel(), yy_normalized.ravel()))
# 预测
Z = clf.predict(coords)
Z = Z.reshape(xx.shape)
# 所有穷举点进行分类展示
light_rgb = ListedColormap(['#AAAAFF', '#FFAAAA', '#AAFFAA'])
plt.pcolormesh(xx, yy, Z, cmap=light_rgb)
plt.scatter(x1, y1, c='b', marker='s', s=50, alpha=0.8)
plt.scatter(x2, y2, c='r', marker='^', s=50, alpha=0.8)
plt.scatter(x3, y3, c='g', s=50, alpha=0.8)
plt.axis((10, 70, 1, 7))
plt.savefig(img.joinpath("KNN分类图.png"))

Z_proba = clf.predict_proba(coords)
Z_proba_reds = Z_proba[:, 1].reshape(xx.shape)
plt.pcolormesh(xx, yy, Z_proba_reds, cmap='Reds')
plt.scatter(x1, y1, c='b', marker='s', s=50, alpha=0.8)
plt.scatter(x2, y2, c='r', marker='^', s=50, alpha=0.8)
plt.scatter(x3, y3, c='g', s=50, alpha=0.8)
plt.axis((10, 70, 1, 7))
plt.savefig(img.joinpath("KNN概率图.png"))
